import time

import torch
from torch_geometric.nn import GATConv,TAGConv
from torch_geometric.data import Data
from torch import nn
  
# 定义 GATConv 模型  
class GATModel(torch.nn.Module):  
    def __init__(self, in_channels, out_channels):  
        super(GATModel, self).__init__()  
        self.conv = TAGConv(in_channels, out_channels)
  
    def forward(self, data):  
        x, edge_index = data.x, data.edge_index  
        x = self.conv(x, edge_index)

        return x  
  
# 创建数据  
edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long)  
x = torch.tensor([[-1], [0], [1], [2]], dtype=torch.float)

train_mask = torch.ones(len(x), dtype=torch.bool)
y_hat = torch.ones(len(x)).long()

y_hat = torch.LongTensor([1,1,0,0])
print(y_hat,"yhat")
data = Data(x=x, edge_index=edge_index,train_mask=train_mask,y=y_hat)
print(data)

# 定义模型和优化器  
model = GATModel(1, 2)  
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
# 训练模型  
for epoch in range(1000):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    print("x_out",out)
    # time.sleep(1)
    # print(out[data.train_mask])
    # print(data.y[data.train_mask])
    loss =  criterion(out[data.train_mask], data.y[data.train_mask])
    # loss = torch.nn.functional.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch}, Loss: {loss.item()}")
